#include<stdio.h>
#include<math.h>
#include<stdlib.h>
#include<time.h>
 
#define PI 3.1415926
 
void DFT(float * input,int Nx,int Nz,float * real,float * imag);
void IDFT(float *real_part,float *imag_part,int Nx,int Nz,float *idft_part);
 
int main()
{
	int it,fm;
	float t1,t0,N,dt,L2;
	clock_t start,end;
	FILE *fp;
	N=1024*1;dt=0.0005;t0=0.1;fm=25;
 
	float *ricker=(float *)malloc(sizeof(float)*N);
     int f1 = 100;//   	# 信号的特征频率1
	int f2 = 300;//   # 信号的特征频率2
	int f3 = 500;//   # 信号的特征频率3
	for(int i=0;i<N;i++){
		//t1=it*dt-t0;
		//ricker[it]=(1.0-2.0*PI*PI*fm*fm*t1*t1)*exp(-(PI*PI*fm*fm*t1*t1));	
		float j = i;
		ricker[i] = 10 + 7 * sin(2*PI*f1*(j*1/N)) + 5 * sin(2*PI*f2* (j*1/N)) + 3 * sin(2*PI*f3* (j*1/N));

	}
	fp=fopen("ricker.txt","wb");
	//fwrite(ricker,sizeof(float),N,fp);
	//fclose(fp);
 	for(int i = 0; i < N ; i++)
		fprintf(fp,"%f,",ricker[i]);
	fclose(fp);

	float *Real_part=(float *)malloc(sizeof(float)*N);
	float *Imag_part=(float *)malloc(sizeof(float)*N);
	float *Idft_part=(float *)malloc(sizeof(float)*N);
	float *residual =(float *)malloc(sizeof(float)*N);
 
	start=clock();
		printf("[%s +%d %s] 1*N=%f\n", __FILE__, __LINE__, __func__, 1*N);

	DFT(ricker,1,N,Real_part,Imag_part);
	end=clock();
	fprintf(stderr,"DFT time taken: %lf\n",(double)((end-start)/CLOCKS_PER_SEC));
 
	start=clock();
	//IDFT(Real_part,Imag_part,1,N,Idft_part);
	end=clock();
	fprintf(stderr,"IDFT time taken: %lf\n",(double)((end-start)/CLOCKS_PER_SEC));
 
	L2=0;
	for(it=0;it<N;it++){
		residual[it]=ricker[it]-Idft_part[it];
		L2+=residual[it]*residual[it];
	}
	fprintf(stderr,"Residual between DFT with IDFT is %f\n",sqrt(L2));
	fp=fopen("ricker_residual.bin","wb");
	fwrite(residual,sizeof(float),N,fp);
	fclose(fp);
	
	free(ricker);
	free(Real_part);
	free(Imag_part);
	free(Idft_part);
    free(residual);
	return 0;
 
}
 
void DFT(float * input,int Nx,int Nz,float * real,float * imag)
{
 
	float *amp=(float *)malloc(sizeof(float)*Nx*Nz);
	float R,I;
	float re,im;
	int ix,in,ik;
	for(ix=0;ix<Nx;ix++)
	{
		for(ik=0;ik<Nz;ik++)
    	{
        		real[ix*Nz+ik]=0;
        		imag[ix*Nz+ik]=0;
 
			R=0;
			I=0;
        	for(in=0;in<Nz;in++)
			{
				R=cos(((-2)*PI*ik*in)/Nz);
				I=sin(((-2)*PI*ik*in)/Nz);
			
				real[ix*Nz+ik] += input[ix*Nz+in]*R;
				imag[ix*Nz+ik] += input[ix*Nz+in]*I;			
        	}
			re=real[ix*Nz+ik];
			im=imag[ix*Nz+ik];
			amp[ix*Nz+ik]=sqrt(pow(re,2)+pow(im,2));
   		}
 
	}
 
	FILE *fp=fopen("ricker_dft_real.txt","wb");
	FILE *fm=fopen("ricker_dft_imag.bin","wb");
	FILE *fr=fopen("ricker_dft_amplitude.txt","wb");
 
	fwrite(real,sizeof(float),Nx*Nz,fp);
	fwrite(imag,sizeof(float),Nx*Nz,fm);
	//fwrite(amp,sizeof(float),Nx*Nz,fr);
	printf("[%s +%d %s] Nx*Nz=%d\n", __FILE__, __LINE__, __func__, Nx*Nz);

  	for(int i = 0; i < Nx*Nz ; i++)
		fprintf(fr,"%f,",amp[i]);

	fclose(fp);
	fclose(fm);
	fclose(fr);
 
	free(amp);
}
 
void IDFT(float *real_part,float *imag_part,int Nx,int Nz,float *idft_part)
{
	float *real=(float *)malloc(sizeof(float)*Nx*Nz);
	float *imag=(float *)malloc(sizeof(float)*Nx*Nz);
	
	float R,I;
	int ix,in,ik;
	for(ix=0;ix<Nx;ix++)
	{
		for(in=0;in<Nz;in++)
    		{
        		real[ix*Nz+in]=0;
        		imag[ix*Nz+in]=0;
 
			R=0;
			I=0;
        		for(ik=0;ik<Nz;ik++)
			{
				R=cos((2*PI*ik*in)/Nz);
				I=sin((2*PI*ik*in)/Nz);
			
				real[ix*Nz+in] += real_part[ix*Nz+ik]*R;
				imag[ix*Nz+in] += imag_part[ix*Nz+ik]*I;			
        		}
			idft_part[ix*Nz+in]=((float)(real[ix*Nz+in]/Nz)-(float)(imag[ix*Nz+in]/Nz));//a-b !!!!
   		}
	}
	
	FILE *fp=fopen("ricker_idft.bin","wb");
	fwrite(idft_part,sizeof(float),Nx*Nz,fp);
 
	fclose(fp);
	free(real);
	free(imag);
}